# modify_request.py
from mitmproxy import http, ctx
import json
from urllib.parse import parse_qs, urlencode
import redis
import re
import time
r = redis.StrictRedis(host='10.8.0.17', port=6379, db=0)
from urllib.parse import urlparse

# 变更授权状态


def update_login_status(user_id, status):
    # 用户散列的key
    user_hash_key = f"user:{user_id}"
    # 如果是新用户或更新用户信息，设置用户散列
    r.hset(user_hash_key, "auth_status", status)


def request(flow: http.HTTPFlow) -> None:
    if "content-autofill.googleapis.com" in flow.request.pretty_url:
        flow.kill()
    # 检查是否是特定的接口
    if "/query/corpadmin/grantUser" in flow.request.pretty_url or "/query/corpadmin/revokeUser" in flow.request.pretty_url:
        while True:
            json_data = r.lpop("user_grant_queue")
            ctx.log.info(f"origin url>>>>>>>>{flow.request.pretty_url}")
            ctx.log.info(f"授权队列>>>>>>>>{json_data}")
            if not json_data:
                break
            else:
                # 反序列化JSON字符串回字典
                grant_info = json.loads(json_data)
                if grant_info["expire_time"] < time.time():
                    continue
                user_id = grant_info["user_id"]
                operation = grant_info["operation"]
                url = flow.request.path
                if operation == "grant":
                    new_url = re.sub(r'revokeUser\.json',
                                     'grantUser.json', url)
                    if new_url != url:
                        flow.request.path = new_url
                        flow.request.query["_api"] = 'CorpAdmin.grantUser'
                elif operation == "revoke":
                    new_url = re.sub(r'grantUser\.json',
                                     'revokeUser.json', url)
                    if new_url != url:
                        flow.request.path = new_url
                        flow.request.query["_api"] = 'CorpAdmin.revokeUser'
                content_type = flow.request.headers.get("Content-Type", "")
                # 如果内容类型是x-www-form-urlencoded，则进行处理
                if "x-www-form-urlencoded" in content_type:
                    # 解析请求体为字典
                    request_form_data = parse_qs(flow.request.get_text())
                    if "workNo" in request_form_data:
                        request_form_data["workNo"] = user_id
                    # 将修改后的数据重新编码成字符串并设置为请求体
                    flow.request.set_text(
                        urlencode(request_form_data, doseq=True))
                break
                # 在这里添加你想要对请求进行的修改
                # 例如: 修改query参数、headers或者请求体等
                # flow.request.query["new_param"] = "value"  # 示例：添加一个新的query参数


def response(flow: http.HTTPFlow) -> None:
    if "/query/corpadmin/grantUser" in flow.request.pretty_url or "/query/corpadmin/revokeUser" in flow.request.pretty_url:
        response = flow.response
        json_data = json.loads(response.get_text())
        request_form_data = parse_qs(flow.request.get_text())
        ctx.log.info(
            f"resonse>>>>>>>>>>>>{json_data}>>>>>>>{request_form_data['workNo']}")
        content_type = flow.request.headers.get("Content-Type", "")
        # 如果内容类型是x-www-form-urlencoded，则进行处理
        if "x-www-form-urlencoded" in content_type:
            # 解析请求体为字典
            request_form_data = parse_qs(flow.request.get_text())
            _api = flow.request.query["_api"]
            if "workNo" in request_form_data:
                user_id = request_form_data["workNo"]
                parsed_url = urlparse(flow.request.pretty_url)
                domain = parsed_url.netloc
                # 分割域名并提取所需部分
                domain_parts = domain.split('.')
                extracted_part = domain_parts[0]
                if user_id:
                    if (isinstance(user_id, list)):
                        user_id = user_id[0]
                    if _api == 'CorpAdmin.grantUser':
                        current_time = time.time()
                        r.zadd(f'{extracted_part}:users_by_last_login', {user_id: current_time})
                    elif _api == 'CorpAdmin.revokeUser':
                        r.zrem(f'{extracted_part}:users_by_last_login', user_id)
        # ctx.log.info(f"URL: {flow.request.pretty_url}")
        # try:
        #     # 获取响应的文本数据
        #     json_data = json.loads(response.get_text())
        #     # 处理或打印JSON数据
        #     ctx.log.info(f"URL: {flow.request.pretty_url}")
        #     ctx.log.info(f"JSON Data: {json_data}")
        #     if json_data["success"] == True:
        #         update_login_status(user_id)
        # except json.JSONDecodeError:
        #     ctx.log.info("Response is not in JSON format")
